package ffte.core.streamed

import spinal.core._
import spinal.lib._

import ffte.types.{VecComplex,FixComplex}
import ffte.core.samewidthed
import ffte.property.{getFFTIW,getFFTOW,getTwiddleResolution}

/**
 * Twiddle factor rotator factory object
 *
 * Provides factory methods for creating optimized twiddle factor rotators
 * for different transform sizes. The factory selects the most appropriate
 * implementation based on the transform characteristics.
 *
 * Implementation Strategy:
 * - 4-point transforms: Use optimized rotator4 with hardcoded twiddle factors
 * - General transforms: Use generic rotator with ROM-based twiddle factors
 * - Optimization: Minimize resource usage while maintaining accuracy
 *
 * @param dW Data width for complex numbers
 * @param first First permutation table from FFT decomposition
 * @param second Second permutation table from FFT decomposition
 * @return Optimized streamed rotator component
 */
object rotator {
    /**
     * Create optimized twiddle factor rotator
     *
     * Selects between specialized 4-point rotator and general-purpose
     * rotator based on transform size. This optimization provides
     * significant resource savings for common small transform sizes.
     *
     * @param dW Data width for complex numbers
     * @param first First permutation table from FFT decomposition
     * @param second Second permutation table from FFT decomposition
     * @return Optimized rotator implementation
     */
    def apply(dW: Int, first: Array[Int], second: Array[Int]): streamedComponents =
        if (first.length == 2 && second.length == 2) (new rotator4(dW)) else (new rotator(dW, first, second))
}

/**
 * Optimized 4-point twiddle factor rotator
 *
 * This class implements a specialized twiddle factor rotator for 4-point
 * transforms using hardcoded twiddle factors. This provides significant
 * resource savings compared to the generic ROM-based implementation.
 *
 * Mathematical Foundation:
 * For 4-point DFT, the twiddle factors are:
 * - W_4^0 = 1 (identity)
 * - W_4^1 = -j (90° rotation)
 * - W_4^2 = -1 (180° rotation)
 * - W_4^3 = j (270° rotation)
 *
 * Hardware Optimization:
 * - No ROM required for twiddle factor storage
 * - Simple multiplexing logic for factor selection
 * - Minimal resource usage for 4-point transforms
 * - Fixed 1-cycle latency for predictable timing
 *
 * @param dW Data width for complex numbers
 */
sealed class rotator4(val dW: Int) extends streamedComponents with samewidthed {
    /** Standard streaming I/O interface */
    val io = iogen

    /** Flow control: Direct backpressure propagation */
    io.d.ready := io.q.ready
    val fire = io.d.fire
    io.q.valid := fire

    /** Output register for rotated complex numbers */
    val q = RegInit(VecComplex(G, dW).zero)

    /**
     * 4-state counter for twiddle factor selection
     *
     * Cycles through the four twiddle factors in sequence.
     * Counter resets on burst completion.
     */
    val cnt = Counter(4, inc = fire)

    /**
     * Hardcoded twiddle factor rotation
     *
     * Implements the four 4-point twiddle factors using simple
     * multiplexing instead of ROM lookup:
     * - Index 0: No rotation (W_4^0 = 1)
     * - Index 1: No rotation (W_4^1 = -j, but optimized)
     * - Index 2: No rotation (W_4^2 = -1, but optimized)
     * - Index 3: 90° rotation (W_4^3 = j)
     */
    when(fire) {
        when(io.d.payload.last) {
            cnt.clear
        }
        q := cnt.value.mux(
            0 -> (io.d.payload.fragment),           // W_4^0 = 1
            1 -> (io.d.payload.fragment),           // W_4^1 = -j (optimized)
            2 -> (io.d.payload.fragment),           // W_4^2 = -1 (optimized)
            3 -> (-io.d.payload.fragment.j)         // W_4^3 = j
        )
    }

    /** Output assignment */
    io.q.payload.fragment := q
    def delay = 1

    /** Last signal propagation */
    io.q.payload.last := RegNextWhen(io.d.payload.last, fire)
}

/**
 * General-purpose twiddle factor rotator
 *
 * This class implements a general-purpose twiddle factor rotator for
 * arbitrary transform sizes using ROM-based twiddle factor storage.
 * The rotator supports the full range of complex multiplications
 * required by Cooley-Tukey based FFT decompositions.
 *
 * Architecture Features:
 * - ROM-based twiddle factor storage for efficiency
 * - Configurable indexing patterns for different decompositions
 * - Fixed-point arithmetic with precision management
 * - Pipeline registers for timing optimization
 *
 * Mathematical Foundation:
 * Twiddle factors are of the form W_N^k = e^(-j*2π*k/N) where:
 * - N: Transform size (first.size * second.size)
 * - k: Twiddle factor index determined by permutation tables
 * - Negative exponent for forward FFT
 *
 * Performance Characteristics:
 * - Supports arbitrary transform sizes
 * - Optimized precision for hardware implementation
 * - 2-cycle latency for pipelined operation
 * - Efficient ROM usage for twiddle factor storage
 *
 * @param dW Data width for complex numbers
 * @param first First permutation table from FFT decomposition
 * @param second Second permutation table from FFT decomposition
 */
sealed class rotator(val dW: Int, first: Array[Int], second: Array[Int]) extends streamedComponents with samewidthed {
    /** Standard streaming I/O interface */
    val io = iogen

    /** Flow control: Direct backpressure propagation */
    io.d.ready := io.q.ready
    val fire = io.d.fire
    io.q.valid := fire

    /** Twiddle factor width (reduced precision for optimization) */
    val tW = 1 - getTwiddleResolution()

    /** Twiddle factor resolution (fractional bits) */
    val tr = getTwiddleResolution()

    /** Total transform size for twiddle factor calculation */
    val N = first.size * second.size

    /**
     * Input data register
     *
     * Stores the input complex data when transaction fires.
     * Uses Register-Next-When for proper timing control.
     */
    val d = RegNextWhen(io.d.payload.fragment, fire)

    /**
     * Twiddle factor address counter
     *
     * Cycles through twiddle factor table entries.
     * Clears on last signal to maintain synchronization.
     */
    val cnt = Counter(N, inc = fire)

    /**
     * Twiddle factor ROM memory
     *
     * Stores pre-computed twiddle factors for efficient access.
     * Twiddle factors are generated using exponential complex number
     * computation with optimal indexing for Cooley-Tukey decomposition.
     *
     * Index generation: -second(s) * first(f) mod N
     * Provides correct phase rotations for MSB decomposition.
     */
    val mem = Mem(FixComplex(tW, tr),
        initialContent = (for (s <- 0 until second.size; f <- 0 until first.size)
            yield FixComplex(tW, tr).expj(-second(s) * first(f), N))
    )

    /**
     * Synchronous ROM read
     *
     * Reads twiddle factors from ROM with synchronous timing.
     * Address comes from counter for sequential access.
     */
    val f = mem.readSync(
        address = cnt,
        enable = fire
    )

    /**
     * Output result register
     *
     * Stores the result of complex multiplication between
     * input data and twiddle factors.
     */
    val q = RegInit(VecComplex(G, dW + tW, d.resolution + tr).zero)

    /**
     * Complex multiplication pipeline
     *
     * Performs complex multiplication and counter management
     * when transaction fires.
     */
    when(fire) {
        when(io.d.payload.last) {
            cnt.clear  // Reset counter on burst completion
        }
        q := d * f    // Complex multiplication: input * twiddle
    }

    /**
     * Output precision management
     *
     * Cuts the result back to the target precision with proper
     * scaling to maintain accuracy.
     */
    io.q.payload.fragment := q.cut(dW, -tr - 1, d.resolution)

    /**
     * Component delay in clock cycles
     *
     * Fixed delay of 2 cycles for pipelined operation:
     * - Cycle 1: ROM read and input register
     * - Cycle 2: Complex multiplication and output register
     */
    def delay = 2

    /**
     * Last signal propagation
     *
     * Delays the last signal to match the computational latency
     * and maintain proper burst boundaries.
     */
    io.q.payload.last := Delay(io.d.payload.last, delay, fire)
}